{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intralayer interatomic distance analysis of LAMMPS MD trajectory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmpdump as lmpdump    # Use lmpdump.py parser\n",
    "from sitator.util import PBCCalculator    # Use PBC calculator from sitator package\n",
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "\n",
    "pwd = os.getcwd()\n",
    "\n",
    "print('Please cite DOI: 10.1021/acs.jpcc.9b04863.')\n",
    "print('This script analyzes intralayer interatomic distances from LAMMPS MD trajectory.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('Note 1: Intended only for periodic slab models of up to bimetallic FCC systems, with vacuum along the z-direction.')\n",
    "print('Note 2: Intended only for LAMMPS NVT simulations.\\n')\n",
    "\n",
    "print('Note 3: Requires a LAMMPS DATA file containing the equilibrated box dimensions.')\n",
    "print('Note 4: Requires a trajectory ordered by atom ID, unscaled, wrapped.')\n",
    "print('Note 5: Requires lmpdump.py to load the LAMMPS trajectories.')\n",
    "print('Note 6: Requires PBC calculator from sitator package.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load LAMMPS trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LAMMPS custom-style dump file - Clamped trajectory\n",
    "name = input('Enter the name of the LAMMPS trajectory file (ID, type, x, y, z) : \\nWarning: Must be ordered by atom ID, unscaled, and wrapped! Please quit now if not so.\\n')\n",
    "print('Loading trajectory...')\n",
    "file = pwd + '/' + name\n",
    "xyz = lmpdump.lmpdump(file, loadmode='all')\n",
    "print('\\n')\n",
    "\n",
    "print('Loading time steps...')\n",
    "step = []\n",
    "for s in xyz.finaldict.keys():    # keys = time steps\n",
    "    step.append(s)\n",
    "print('Done.\\n')\n",
    "\n",
    "size = np.array(step).size    # Number of time steps\n",
    "N_dump = xyz.finaldict[step[0]][1].shape[0]    # Number of atoms dumped\n",
    "xlo_ext = xyz.finaldict[step[0]][0][0]    # x_ext = x + xy\n",
    "xhi_ext = xyz.finaldict[step[0]][0][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract unit cell information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load LAMMPS data file after NPT equilibration\n",
    "name = input('Enter the name of the LAMMPS DATA file containing the equilibrated box dimensions : ')\n",
    "file_eq = pwd + '/' + name\n",
    "log = open(file_eq,'r')\n",
    "log_lines = log.readlines()\n",
    "log_lines = [line.split() for line in log_lines]\n",
    "\n",
    "for line_index, line in enumerate(log_lines):\n",
    "    if line:\n",
    "        for l in line:\n",
    "            \n",
    "            if l == 'atoms':\n",
    "                N_atom = int(line[0])\n",
    "                \n",
    "            elif l == 'types':\n",
    "                N_type = int(line[0])\n",
    "\n",
    "            elif l == 'xlo':\n",
    "                xlo = float(line[0])\n",
    "                xhi = float(line[1])\n",
    "            \n",
    "            elif l == 'ylo':\n",
    "                ylo = float(line[0])\n",
    "                yhi = float(line[1])\n",
    "            \n",
    "            elif l == 'zlo':\n",
    "                zlo = float(line[0])\n",
    "                zhi = float(line[1])\n",
    "            \n",
    "            elif l == 'xy':\n",
    "                xy = float(line[0])\n",
    "                xz = float(line[1])\n",
    "                yz = float(line[2])\n",
    "            \n",
    "            elif l == 'Atoms':\n",
    "                head = line_index + 1\n",
    "                start = line_index + 2\n",
    "            \n",
    "            elif l == 'Velocities':\n",
    "                end = line_index - 1\n",
    "\n",
    "Lx = xhi - xlo    # Box dimensions\n",
    "Ly = yhi - ylo\n",
    "Lz = zhi - zlo\n",
    "tan = Ly/xy\n",
    "\n",
    "lat = np.array([[xhi-xlo, 0.0, 0.0], [xy, yhi-ylo, 0.0], [xz, yz, zhi-zlo]])    # Lattice matrix\n",
    "pbc = PBCCalculator(lat)    # Periodic boundary condition\n",
    "\n",
    "ID2 = int(input('Enter the ID of a nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID3 = int(input('Enter the ID of another nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID4 = int(input('Enter the ID of an atom that lies one layer above atom #1 : '))\n",
    "\n",
    "for line in log_lines[start:end]:\n",
    "\n",
    "    line[:2] = np.array(line[:2]).astype(int)    # Atom ID, atom type\n",
    "    line[2:5] = np.array(line[2:5]).astype(float)    # x, y, z\n",
    "    line[5:8] = np.array(line[5:8]).astype(int)    # ix, iy, iz (inverse sign)\n",
    "            \n",
    "    x = line[2] + line[5]*Lx + line[6]*xy    # Wrapped coordinates\n",
    "    y = line[3] + line[6]*Ly\n",
    "    z = line[4] + line[7]*Lz\n",
    "    \n",
    "    if line[0] == 1:\n",
    "        atom1 = np.array([x,y,z])\n",
    "    \n",
    "    elif line[0] == ID2:\n",
    "        atom2 = np.array([x,y,z])\n",
    "    \n",
    "    elif line[0] == ID3:\n",
    "        atom3 = np.array([x,y,z])\n",
    "    \n",
    "    elif line[0] == ID4:\n",
    "        atom4 = np.array([x,y,z])\n",
    "\n",
    "r12 = np.linalg.norm(np.subtract(atom1,atom2))\n",
    "r13 = np.linalg.norm(np.subtract(atom1,atom3))\n",
    "dr_avg = np.mean([r12, r13])    # In-plane NN distance\n",
    "\n",
    "a1 = np.subtract(atom2,atom1)    # 1st intralayer vector\n",
    "a1 = a1 / np.linalg.norm(a1)\n",
    "\n",
    "a2 = np.subtract(atom3,atom1)    # 2nd intralayer vector\n",
    "a2 = a2 / np.linalg.norm(a2)\n",
    "\n",
    "a3 = np.cross(a1,a2)    # Normal vector\n",
    "if a3[2] < 0:\n",
    "    a3 = -a3    # Ensure normal vector along +z\n",
    "a3 = a3 / np.linalg.norm(a3)\n",
    "\n",
    "dn_avg = np.dot(atom4,a3) - np.dot(atom1,a3)    # Interlayer distance; n = normal component\n",
    "\n",
    "N_slab = int(input('Enter the number of layers in the slab, excluding the deposit/adsorbate layer(s) : '))\n",
    "if xyz.finaldict[step[0]][1]['id'][0] != 1:    # Is bottommost layer muted?\n",
    "    N_slab -= 1\n",
    "    \n",
    "z_atom = []    # Atomic height\n",
    "for atom in range(0,N_dump):\n",
    "    xs = xyz.finaldict[step[-1]][1]['xs'][atom]    # Use last frame\n",
    "    ys = xyz.finaldict[step[-1]][1]['ys'][atom]\n",
    "    zs = xyz.finaldict[step[-1]][1]['zs'][atom]\n",
    "    rs = np.array([xs,ys,zs])\n",
    "    r = np.dot(lat,rs)\n",
    "    z_atom.append(r[2])\n",
    "z_atom = np.array(z_atom)\n",
    "z_atom = np.sort(z_atom)    # Sort by ascending height\n",
    "\n",
    "atom_next = 0\n",
    "z_layer = []    # Average layer heights\n",
    "while atom_next < N_dump-1:\n",
    "    z_bin = []\n",
    "    \n",
    "    for atom in range(atom_next,N_dump):\n",
    "        z0 = z_atom[atom]    # Use last frame\n",
    "        z1 = z_atom[atom+1]\n",
    "        dz = z1 - z0\n",
    "        \n",
    "        if abs(dz) > 0.20*dn_avg:\n",
    "            atom_next = atom+1\n",
    "            break\n",
    "            \n",
    "        elif z0 not in z_bin:\n",
    "            z_bin.append(z0)\n",
    "        \n",
    "        if atom == N_dump-2:\n",
    "            atom_next = atom+1\n",
    "            break\n",
    "        \n",
    "    z_bin = np.array(z_bin)\n",
    "    z_avg = np.mean(z_bin)\n",
    "    z_layer.append(z_avg)\n",
    "    \n",
    "N_layer = len(z_layer)\n",
    "\n",
    "log.close()\n",
    "\n",
    "print('Unit cell information & normal direction extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intralayer interatomic distance analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_ravg = pwd + '/ravg.txt'\n",
    "out_ravg = open(file_ravg,'w')\n",
    "\n",
    "header = 'time (ns)\\t'\n",
    "for l in range(0,N_layer):\n",
    "    name = 'L' + str(l)\n",
    "    header += name + '\\t'\n",
    "header += '\\n'\n",
    "out_ravg.write(header)\n",
    "\n",
    "dt = int(input('Enter the simulation timestep in fs : '))\n",
    "\n",
    "for s_index, s in enumerate(tqdm(step)):\n",
    "    t = s*dt/(10**6)\n",
    "    out = str(t) + '\\t'\n",
    "    \n",
    "    layer = {}    # Sort atoms by layer\n",
    "    for l in range(0,N_layer):\n",
    "        name = 'L' + str(l)\n",
    "        layer[name] = []\n",
    "    \n",
    "    for atom in range(0,N_dump):\n",
    "        \n",
    "        xs = xyz.finaldict[s][1]['xs'][atom]\n",
    "        ys = xyz.finaldict[s][1]['ys'][atom]\n",
    "        zs = xyz.finaldict[s][1]['zs'][atom]\n",
    "        rs = np.array([xs,ys,zs])\n",
    "        \n",
    "        r = np.dot(lat,rs)\n",
    "        z = r[2]\n",
    "        \n",
    "        for l in range(0,N_layer):\n",
    "            name = 'L' + str(l)\n",
    "            z_ref = z_layer[l]\n",
    "            dz = z - z_ref\n",
    "            if abs(dz) < 0.20*dn_avg:\n",
    "                layer[name].append(r)\n",
    "                break\n",
    "    \n",
    "    for l in range(0,N_layer):\n",
    "        name = 'L' + str(l)\n",
    "        layer[name] = np.array(layer[name])\n",
    "        \n",
    "        N = layer[name].shape[0]\n",
    "        if N <= 1:\n",
    "            rij_avg = 0.0\n",
    "            out += str(rij_avg) + '\\t'\n",
    "            continue\n",
    "        \n",
    "        rij = pbc.pairwise_distances(layer[name])    # Pairwise distance matrix        \n",
    "        rij_bin = []\n",
    "        for i in range(0,N):\n",
    "            for j in range(0,N):\n",
    "                if j > i:\n",
    "                    if rij[i][j] < 1.20*dr_avg:\n",
    "                        rij_bin.append(rij[i][j])                        \n",
    "        rij_bin = np.array(rij_bin)\n",
    "\n",
    "        if rij_bin.shape[0] == 0:\n",
    "            rij_avg = 0.0\n",
    "        else:\n",
    "            rij_avg = np.mean(rij_bin)\n",
    "        \n",
    "        out += str(rij_avg) + '\\t'\n",
    "    \n",
    "    out += '\\n'\n",
    "    out_ravg.write(out)\n",
    "    \n",
    "out_ravg.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
